from mesa.space import MultiGrid
from mesa import Model, Agent
from mesa.time import RandomActivation
import numpy as np
from mesa.visualization.UserParam import UserSettableParameter

from mesa.visualization.modules import CanvasGrid
from mesa.visualization.ModularVisualization import ModularServer


class GridAgent(Agent):

    def __init__(self, unique_id, model, stat: str):
        super().__init__(unique_id, model)
        self.wealth = 1
        self.stat = stat  # 有三种：blank\tree\flame.
        self.flameLife = 1  # 火焰只持续一个刷新周期

    def walk(self):  # 相当于火的扩散
        surroundings = self.model.grid.get_neighborhood(
            self.pos,
            moore=True,
            include_center=False)
        if (self.stat == 'flame'):
            if (self.flameLife > 0):
                self.flameLife -= 1
            else:
                self.stat = 'blank'
        surroundings = [self.getElementPos(1, 0), self.getElementPos(-1, 0), self.getElementPos(0, 1),
                        self.getElementPos(0, -1)]  # 在本程序的设计中，火焰只会向上下左右四个邻居扩散，不向对角的单元格扩散。

        cellmate = self.model.grid.get_cell_list_contents(surroundings)

        for mate in cellmate:
            if (mate.stat == 'flame') & (self.stat == 'tree'):
                self.stat = 'flame'

    def move(self, dx: int, dy: int):
        new_position = (self.pos[0] + dx, self.pos[1] + dy)  # 向右为正，向上为正。
        self.model.grid.move_agent(self, new_position)

    def getElementPos(self, dx: int, dy: int):  # 输入相对位置返回绝对位置
        x, y = self.pos[0] + dx, self.pos[1] + dy
        width = self.model.grid.width
        print(self.pos, dx, dy, x, y, width)
        if (x >= width):
            x -= width
        elif x < 0:
            x += width

        if (y >= width):
            y -= width
        elif y < 0:
            y += width

        return (x, y)

    def give_money(self):
        cellmates = self.model.grid.get_cell_list_contents([self.pos])
        cellmatesCross = [self.getElementPos(1, 0), self.getElementPos(-1, 0), self.getElementPos(0, 1),
                          self.getElementPos(0, -1)]
        if len(cellmates) > 1:
            other = self.random.choice(cellmates)
            other.wealth += 1
            self.wealth -= 1

    def step(self):
        self.walk()


class GridModel(Model):
    """A model with some number of agents."""

    def __init__(self, treeDensity, width, height):

        self.grid = MultiGrid(width, height, True)
        self.schedule = RandomActivation(self)
        self.treeDensity = treeDensity

        for x in range(self.grid.width):
            for y in range(self.grid.height):

                if (np.random.rand() < self.treeDensity):
                    if (x <= 1):
                        a = GridAgent(x * self.grid.width + y, self, 'flame')
                    else:
                        a = GridAgent(x * self.grid.width + y, self, 'tree')
                else:
                    a = GridAgent(x * self.grid.width + y, self, 'blank')
                self.schedule.add(a)
                self.grid.place_agent(a, (x, y))

    def step(self):
        self.schedule.step()

    def running(self):
        self.step()


def agent_portrayal(agent):
    portrayal = {"Shape": "rect",  # 或者 "circle"等等各种类型
                 "Filled": "true",
                 "w": 1, "h": 1}

    if agent.stat == 'flame':
        portrayal["Color"] = "red"
        portrayal["Layer"] = 1
        portrayal["w"] = 1
        portrayal['h'] = 1
    elif agent.stat == 'tree':
        portrayal["Color"] = "green"
        portrayal["Layer"] = 1
        portrayal["w"] = 1
        portrayal['h'] = 1
    elif agent.stat == 'blank':
        portrayal["Color"] = "black"
        portrayal["Layer"] = 1
    return portrayal


width = 30
height = 30

model_params = {

    "height": height,
    "width": width,
    "treeDensity": UserSettableParameter("slider", "森林密度", 0.6, 0.01, 1, 0.01)
}
grid = CanvasGrid(agent_portrayal, width, height, 500, 500)  # 像素范围500*500，有width*height个网格

server = ModularServer(GridModel,
                       [grid],
                       "Money Model",
                       model_params=model_params)
server.port = 8521  # 设定端口。如果做了内网穿透，甚至可以远程访问。
server.launch()
